
from PIL import Image

import cv2
import numpy as np
import matplotlib.pyplot as plt


def split_channels(img):
    num_channels = len(img.shape)
    result = []
    # 按 dim=2 切分，得到 (height, width, 1)
    channels = np.split(img, num_channels, axis=2)
    for c in channels:
        # 按 dim=2 求和，得到 (height, width)
        out = np.sum(c, axis=2, keepdims=False)
        result.append(out)
    return result
    # or
    # return img[:, :, 0], img[:, :, 1], img[:, :, 2]


def join_channels(channels):
    # (height, width, 1) -> (height, width, 3)
    return np.concatenate(channels, axis=2)


def from_rgb(args):
    if args.tool == 'cv2':
        img = cv2.imread(args.src, cv2.IMREAD_COLOR)
        # (height, width, channels=3,BGR) by default
        print(img.dtype, img.shape)
        # img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        # # (height, width, channels=3,RGB)
        # print(img.dtype, img.shape)
    elif args.tool == 'PIL':
        img = Image.open(args.src)
        # (width, height)
        # print(img.size)
        img = np.array(img)
        # (height, width, channels=3,RGB)
        print(img.dtype, img.shape)

    # 放缩
    # resize_wh = (640, 360)
    # img = cv2.resize(img, resize_wh)
    # # (height=360, width=640, channels=3)
    # print(img.shape)

    if args.dtype is not None:
        if args.dtype == 'float32':
            if (img.dtype == np.uint8):
                # to [0, 1]
                img = img / 255.0
            img = img.astype(np.float32)
            print(img.dtype, img.shape)
        else:
            raise ValueError('Unknown dtype: %s' % args.dtype)

    cmap = None
    if args.transform is not None:
        if args.transform == 'swap-y':
            # 上下翻转
            img = img[::-1, :]
        elif args.transform == 'swap-x':
            # 左右翻转
            img = img[:, ::-1]
        elif args.transform == 'swap-xy':
            # 上下左右翻转
            img = img[::-1, ::-1]
        elif args.transform == 'swap-w':
            # 宽高互换
            img = img.transpose(1, 0, 2)
        elif args.transform == 'rgb':
            # 本来就是 RGB
            pass
        elif args.transform == 'bgr':
            # RGB -> BGR
            tmp = img.copy()
            # img[:, :, 0], img[:, :, 2] = tmp[:, :, 2], tmp[:, :, 0]
            img[:, :, (0, 2)] = tmp[:, :, (2, 0)]
        elif args.transform == 'r':
            # 抹掉 G B 通道
            r = img[:, :, 0]
            img = cv2.merge([r, np.zeros(r.shape, np.uint8), np.zeros(r.shape, np.uint8)])
            # or
            # img = np.dstack((r, np.zeros(r.shape, np.uint8), np.zeros(r.shape, np.uint8)))
            # or
            # img[:, :, (1, 2)] = 0
        elif args.transform == 'g':
            # 抹掉 R B 通道
            img[:, :, (0, 2)] = 0
        elif args.transform == 'b':
            # 抹掉 R G 通道
            img[:, :, (0, 1)] = 0
        elif args.transform == 'ndvi':
            # 归一化植被指数 NDIV = (NIR - RED) / (NIR + RED)
            nir = img[:, :, 0]
            red = img[:, :, 2]
            ndvi = (nir - red) / (nir + red)
            # cmap = 'terrain_r'
            # float64 (height, width)
            print(ndvi.dtype, ndvi.shape)
            # to [0, 1]
            ndvi = (ndvi + 1) / 2.0
            # float64 (height, width)
            print(ndvi.dtype, ndvi.shape)
            img = np.repeat(ndvi[:, :, np.newaxis], repeats=3, axis=2)
            # float64 (height, width, channels=3)
            print(img.dtype, img.shape)
        elif args.transform == 'mean':
            m = np.mean(img / 255.0, axis=2)
            # float64 (height, width)
            print(m.dtype, m.shape)
            # img = cv2.merge([m, m, m])
            # or
            # mmm = (m,)*3
            # img = np.stack(mmm, axis=-1)
            # or
            # img = np.repeat(np.expand_dims(m, axis=2), 3, axis=2)
            # or
            img = np.repeat(m[:, :, np.newaxis], repeats=3, axis=2)
            # float64 (height, width, channels=3)
            print(img.dtype, img.shape)
        elif args.transform == 'gray-mean':
            img = np.mean(img, axis=2)
            # (height, width)
            print(img.shape)
            cmap = 'gray'
        elif args.transform == 'gray-L':
            # ITU-R 601-2 亮度转换
            # uint8 -> float64
            img = 0.299 * img[:, :, 0] + 0.587 * img[:, :, 1] + 0.114 * img[:, :, 2]
            # float64 (height, width)
            print(img.dtype, img.shape)
            cmap = 'gray'
        elif args.transform == 'gray-cv2':
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            cmap = 'gray'
        else:
            raise ValueError('Unknown transform: %s' % args.transform)

    # 灰度图的第三维不重要
    if False:
        img = np.expand_dims(img, axis=2)
        # (height, width, channels)
        print(img.shape)

    if args.draw:
        # requires [0..1] for floats or [0..255] for uint8
        plt.imshow(img, cmap=cmap)
        # plt.show()

    if args.dst:
        if args.tool == 'cv2':
            cv2.imwrite(args.dst, img)
        elif args.tool == 'PIL':
            img = Image.fromarray(img)
            img.save(args.dst)


def from_single(args):
    img = cv2.imread(args.src)
    # (height, width, channels)
    print(img.shape)

    plt.imshow(img)
    plt.show()


if __name__ == '__main__':
    import argparse
    mode = 'from_rgb'
    if mode == 'from_rgb':
        parser = argparse.ArgumentParser()
        parser.add_argument('src')
        parser.add_argument('--dst')
        parser.add_argument('--draw', action='store_true')
        parser.add_argument('--dtype')
        parser.add_argument('--tool', default='cv2')
        parser.add_argument('--transform')
        args = parser.parse_args()
        from_rgb(args)
    elif mode == 'from_single':
        parser = argparse.ArgumentParser()
        parser.add_argument('src')
        args = parser.parse_args()
        from_single(args)
